import os
import numpy as np
import scipy.io as sio
from tqdm import tqdm
from sklearn.model_selection import train_test_split

def load_cwru_data(data_path):
    """加载CWRU数据，提取包含DE的信号"""
    data_files = [f for f in os.listdir(data_path) if f.endswith('.mat')]
    processed_data = {}
    
    for file in tqdm(data_files, desc="处理数据文件"):
        file_path = os.path.join(data_path, file)
        mat_data = sio.loadmat(file_path)
        
        # 获取故障类别（从文件名中提取）
        if file == 'normal.mat':
            fault_type = 'normal'
        else:
            # 从文件名中提取故障类型（例如：real_ball_07_118.mat -> ball）
            parts = file.split('_')
            fault_type = parts[1]  # 获取故障类型（ball/inner/outer）
            fault_size = parts[2]  # 获取故障尺寸（07/14/21）
            fault_type = f"{fault_type}_{fault_size}"  # 组合为 ball_07, inner_14 等
        
        if fault_type not in processed_data:
            processed_data[fault_type] = []
        
        # 查找包含DE的变量
        de_keys = [key for key in mat_data.keys() if 'DE' in key and key not in ['__header__', '__version__', '__globals__']]
        
        if not de_keys:
            print(f"警告：在文件 {file} 中未找到包含DE的变量")
            continue
            
        # 使用找到的第一个DE变量
        de_signal = mat_data[de_keys[0]].flatten()
        print(f"文件 {file} 使用变量: {de_keys[0]}")
        
        # 将长信号切分为1200长度的片段
        segment_length = 1200
        for i in range(0, len(de_signal) - segment_length + 1, segment_length):
            segment = de_signal[i:i + segment_length]
            if len(segment) == segment_length:  # 确保片段长度正确
                processed_data[fault_type].append(segment)
    
    return processed_data

def save_processed_data(processed_data, save_path):
    """保存处理后的数据，同时划分训练测试集"""
    os.makedirs(save_path, exist_ok=True)
    
    # 创建训练集和测试集目录
    train_dir = os.path.join(save_path, "train")
    test_dir = os.path.join(save_path, "test")
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)
    
    for fault_type, segments in processed_data.items():
        # 将数据转换为numpy数组
        data_array = np.array(segments)
        
        # 划分训练集和测试集 (80/20)
        train_data, test_data = train_test_split(
            data_array, test_size=0.2, random_state=42
        )
        
        # 分别保存训练集和测试集
        train_file = os.path.join(train_dir, f"{fault_type}.npy")
        test_file = os.path.join(test_dir, f"{fault_type}.npy")
        
        np.save(train_file, train_data)
        np.save(test_file, test_data)
        
        print(f"已保存 {fault_type} 类别: 训练集 {train_data.shape}, 测试集 {test_data.shape}")

def main():
    # 设置路径
    data_path = "./cwru_data"  # 原始数据路径
    save_path = "./cwru_prepro"  # 预处理后数据保存路径
    
    # 处理数据
    print("开始处理CWRU数据...")
    processed_data = load_cwru_data(data_path)
    
    # 保存处理后的数据（包含训练/测试集划分）
    print("保存处理后的数据...")
    save_processed_data(processed_data, save_path)
    print("数据预处理完成！")

if __name__ == "__main__":
    main() 